﻿using Patterns4Net.Attributes;

namespace Patterns4Net.Core.CecilExtensions
{
    using System;
    using System.Collections.Generic;
    using System.Diagnostics.Contracts;
    using System.Linq;
    using System.Reflection;
    using System.Security;
    using Mono.Cecil;
    using Mono.Cecil.Cil;
    using Patterns4Net.Attributes.DesignPatterns;

    public static class StackObserver
    {
        private static IDictionary<OpCode, IOpCodeHandler> handlers = null;

        [ContractClass(typeof(OpCodeHandlerContract))]
        [Strategy]
        private interface IOpCodeHandler
        {
            IEnumerable<OpCode> OpCodesToHandle { get; }

            StackState Handle(StackState previous, Instruction instruction, MethodDefinition methodDefinition);
        }

        /// <summary>
        /// Gets the instances of all handlers.
        /// </summary>
        private static IDictionary<OpCode, IOpCodeHandler> OpCodesHandlers
        {
            get
            {
                Contract.Ensures(Contract.Result<IDictionary<OpCode, IOpCodeHandler>>() != null);
                if (handlers == null)
                {
                    handlers = new Dictionary<OpCode, IOpCodeHandler>();
                    foreach (var instance in new IOpCodeHandler[] { new CallHandler(), new RetHandler(), new NewobjHandler(), new EmptyStackHandler() })
                    {
                        foreach (var opCode in instance.OpCodesToHandle)
                        {
                            Contract.Assert(
                                handlers.ContainsKey(opCode) == false,
                                string.Format("There is already handler for instruction with opCode {0}.", opCode));
                            handlers.Add(opCode, instance);
                        }
                    }
                }

                return handlers;
            }
        }

        /// <summary>
        /// Creates the array of instances of <see cref="StackState"/> (which represents the state of stack at given point) 
        /// for each instruction in given collection.
        /// </summary>
        public static StackState[] CreateStackStates(MethodDefinition methodDefinition)
        {
            Contract.Requires(methodDefinition != null);
            Contract.Requires(methodDefinition.Body != null);
            Contract.Requires(methodDefinition.Body.Instructions != null);
            Contract.Ensures(Contract.Result<IEnumerable<StackState>>() != null);

            var methodInstructions = methodDefinition.Body.Instructions;
            var result = new List<StackState>();
            var previous = new StackState(Enumerable.Empty<StackItem>());
            var defaultHandler = new DefaultHandler();

            foreach (var instruction in methodInstructions)
            {
                IOpCodeHandler handler;
                OpCodesHandlers.TryGetValue(instruction.OpCode, out handler);
                if (handler == null)
                {
                    handler = defaultHandler;
                }

                previous = handler.Handle(previous, instruction, methodDefinition);
                result.Add(previous);
            }

            return result.ToArray();
        }

        // TODO: rename, update comments.
        /// <summary>
        /// For a local variable of reference type it calculates all possible sources of its value.
        /// These may be method's argument, (static) field or new operator call.
        /// If one of the sources is another local variable x, recursively sources of x are added to the result.
        /// </summary>
        public static IEnumerable<StackItem> GetPossibleRoots(MethodDefinition method, StackState[] stackStates, StackItem stackItem)
        {
            if (stackItem.IsLocalVar == false)
            {
                return new[] { stackItem };
            }

            var result = new List<StackItem>();
            int variableIndex = stackItem.LocalVarIndex;
            for (int i = 0; i < method.Body.Instructions.Count; i++)
            {
                var currentInstruction = method.Body.Instructions[i].Normalize(method.Body);
                if (currentInstruction.OpCode == OpCodes.Stloc &&
                    ((VariableReference)currentInstruction.Operand).Index == variableIndex)
                {
                    // the stack item that was assigned to the variable
                    var pushedStackItem = stackStates[i - 1].Last();
                    result.AddRange(GetPossibleRoots(method, stackStates, pushedStackItem));
                }
            }

            return result;
        }
        
        public class Exception : ApplicationException
        {
            public Exception(string methodName, string message, System.Exception inner = null)
                : base(
                    string.Format("Method {0} defined in static class StackObserver has thrown an exception. {1}", methodName, message),
                    inner)
            {                
            }
        }

        [ContractClassFor(typeof(IOpCodeHandler))]
        private abstract class OpCodeHandlerContract : IOpCodeHandler
        {
            public IEnumerable<OpCode> OpCodesToHandle
            {
                get
                {
                    Contract.Ensures(Contract.Result<IEnumerable<OpCode>>() != null);
                    return null;
                }
            }

            public StackState Handle(StackState previous, Instruction instruction, MethodDefinition methodDefinition)
            {
                Contract.Requires(previous != null);
                Contract.Requires(instruction != null);
                Contract.Requires(methodDefinition != null);
                Contract.Ensures(Contract.Result<IEnumerable<OpCode>>() != null);                
                return null;
            }
        }

        /// <summary>
        /// Uses properties <see cref="OpCode.StackBehaviourPush"/> and <see cref="OpCode.StackBehaviourPop"/>.
        /// This class does not implement <see cref="OpCodesToHandle"/> and is not registered as regular OpCode handler.
        /// </summary>
        public class DefaultHandler : IOpCodeHandler
        {
            public IEnumerable<OpCode> OpCodesToHandle
            {
                get
                {
                    throw new NotImplementedException();
                }
            }

            public StackState Handle(StackState previous, Instruction instruction, MethodDefinition methodDefinition)
            {
                var result = previous.Pop(GetPopItemsCount(instruction.OpCode.StackBehaviourPop));
                int pushedItemsCount = GetPushItemsCount(instruction.OpCode.StackBehaviourPush);
                for (int i = 0; i < pushedItemsCount; i++)
                {
                    result = result.Push(new StackItem(methodDefinition, previous, instruction, i));
                }

                return result;
            }

            private static int GetPopItemsCount(StackBehaviour popBehavior)
            {
                switch (popBehavior)
                {
                    case StackBehaviour.Pop0:
                        return 0;
                    case StackBehaviour.Pop1:
                    case StackBehaviour.Popi:
                    case StackBehaviour.Popref:
                        return 1;
                    case StackBehaviour.Pop1_pop1:
                    case StackBehaviour.Popi_pop1:
                    case StackBehaviour.Popi_popi:
                    case StackBehaviour.Popi_popi8:
                    case StackBehaviour.Popi_popr4:
                    case StackBehaviour.Popi_popr8:
                    case StackBehaviour.Popref_pop1:
                    case StackBehaviour.Popref_popi:
                        return 2;
                    case StackBehaviour.Popref_popi_popi:
                    case StackBehaviour.Popref_popi_popi8:
                    case StackBehaviour.Popref_popi_popr4:
                    case StackBehaviour.Popref_popi_popr8:
                    case StackBehaviour.Popref_popi_popref:
                    case StackBehaviour.Popi_popi_popi:
                        return 3;
                    default:
                        // Contract.Assert(false, string.Format("StackBehavior {0} is not supported in this case.", popBehavior));
                        return 0;
                }
            }

            private static int GetPushItemsCount(StackBehaviour pushBehavior)
            {
                switch (pushBehavior)
                {
                    case StackBehaviour.Push0:
                        return 0;
                    case StackBehaviour.Push1:
                    case StackBehaviour.Pushi:
                    case StackBehaviour.Pushr4:
                    case StackBehaviour.Pushr8:
                    case StackBehaviour.Pushi8:
                    case StackBehaviour.Pushref:
                        return 1;
                    case StackBehaviour.Push1_push1:
                        return 2;
                    default:
                        Contract.Assert(false, string.Format("StackBehaviour {0} is not supported in this case.", pushBehavior));
                        return 0;
                }
            }
        }

        /// <summary>
        /// Handles difference between call to void and call to non-void method 
        /// and pops all method arguments.
        /// </summary>
        private class CallHandler : IOpCodeHandler
        {
            public IEnumerable<OpCode> OpCodesToHandle
            {
                get
                {
                    return new[] { OpCodes.Call, OpCodes.Calli, OpCodes.Callvirt };
                }
            }

            public StackState Handle(StackState previous, Instruction instruction, MethodDefinition methodDefinition)
            {
                var method = instruction.Operand as MethodReference;
                Contract.Assert(
                    method != null, 
                    "Call, calli or callvirst instruction does not contain MethodReference instance as operand in StackObserver.CallHandler");
                Contract.Assert(method.MethodReturnType != null);

                StackState result;
                if (method.HasThis)
                {
                    result = previous.Pop(method.Parameters.Count + 1);
                }
                else
                {
                    result = previous.Pop(method.Parameters.Count);
                }

                if (method.MethodReturnType.IsVoid())
                {
                    return result;
                }
                else
                {
                    return result.Push(new StackItem(methodDefinition, previous, instruction));
                }
            }
        }

        /// <summary>
        /// Pops all constructor arguments.
        /// </summary>
        private class NewobjHandler : IOpCodeHandler
        {
            public IEnumerable<OpCode> OpCodesToHandle
            {
                get
                {
                    return new[] { OpCodes.Newobj };
                }
            }

            public StackState Handle(StackState previous, Instruction instruction, MethodDefinition methodDefinition)
            {
                var ctor = ((MethodReference)instruction.Operand);
                return previous.Pop(ctor.Parameters.Count)
                    .Push(new StackItem(methodDefinition, previous, instruction));
            }
        }

        /// <summary>
        /// Handles difference between ret in void method and ret in method with return value.
        /// </summary>
        private class RetHandler : IOpCodeHandler
        {
            public IEnumerable<OpCode> OpCodesToHandle
            {
                get
                {
                    return new[] { OpCodes.Ret };
                }
            }

            public StackState Handle(StackState previous, Instruction instruction, MethodDefinition methodDefinition)
            {
                Contract.Assert(methodDefinition.MethodReturnType != null);

                if (methodDefinition.MethodReturnType.IsVoid())
                {
                    return previous;
                }

                return previous.Pop();
            }
        }

        /// <summary>
        /// Handler for instructions that empty the evaluation stack like leave instruction.
        /// </summary>
        private class EmptyStackHandler : IOpCodeHandler
        {
            public IEnumerable<OpCode> OpCodesToHandle
            {
                get
                {
                    return new[] { OpCodes.Leave, OpCodes.Leave_S };
                }
            }

            public StackState Handle(StackState previous, Instruction instruction, MethodDefinition methodDefinition)
            {
                return previous.Pop(previous.Count);
            }
        }
    }
}
